'''
利用现有模型，针对自己需要解决的问题进行模型的修改
vgg16用来做图片分类可以分1000类，模型输出1000个数据
'''
import torchvision.models
from torch import nn
from torch.utils.data import DataLoader

# vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)

print(vgg16_true)
# train_data = torchvision.datasets.CIFAR10("./dataset", train=True, download=False,
#                                           transform=torchvision.transforms.ToTensor())
# train_dataloader = DataLoader(train_data, batch_size=1, shuffle=True, num_workers=0, drop_last=False)

print("-------")
# vgg16_true.add_module('add_linear', nn.Linear(in_features=1000, out_features=10))
# vgg16_true.classifier.add_module('add_linear', nn.Linear(in_features=1000, out_features=10))

vgg16_true.classifier[6] = nn.Linear(in_features=4096, out_features=10)
print(vgg16_true)
print("------------")

alexnet = torchvision.models.alexnet(pretrained=True)
print(alexnet)
